import os
import numpy as np
import sys
import cv2
import matplotlib.pyplot as plt
from .deepflux.utils import flux_to_skl
# import caffe
sys.path.insert(0, '/disk/liuzhe/PyCharmProjects/skeleton/caffe/python')
import caffe as caffe_fsds

PROTOCOL = {
    'deepflux': '/disk/yami/project/DeepFlux/models/deploy.prototxt',
    'fsds': '/disk/liuzhe/PyCharmProjects/skeleton/models/fsds_test.prototxt',
    'hifi-sklarge': '/disk/yami/project/WSA/models/h2_sklarge_test.prototxt',
    'hifi': '/disk/liuzhe/PyCharmProjects/skeleton/models/h2_test.prototxt',
}

MODEL_PATH = {
    'deepflux': {
        'sklarge': '/disk/yami/project/DeepFlux/models/sklarge_iter_40000.caffemodel',
        'sympascal': '/disk/yami/project/DeepFlux/models/sympascal_iter_40000.caffemodel',
        'symmax300': '/disk/yami/project/DeepFlux/models/SYMMAX300_iter_40000.caffemodel',
        'wh-symmax': '/disk/yami/project/DeepFlux/models/whsymmax_iter_40000.caffemodel',
    },
    'fsds': {
        'sklarge': '/disk/liuzhe/PyCharmProjects/skeleton/snapshot/fsds-sklarge.caffemodel',
        'sympascal': '/disk/liuzhe/PyCharmProjects/skeleton/snapshot/fsds-sympascal_iter_40000.caffemodel',
        'symmax300': '/disk/liuzhe/PyCharmProjects/skeleton/snapshot/fsds-symmax_iter_40000.caffemodel',
        'wh-symmax': '/disk/liuzhe/PyCharmProjects/skeleton/snapshot/fsds-wh-symmax.caffemodel',
    },
    'hifi': {
        'sklarge': '/disk/liuzhe/PyCharmProjects/skeleton/snapshot/h2-skl_iter_40000.caffemodel',
        'sympascal': '/disk/liuzhe/PyCharmProjects/skeleton/snapshot/h2-sympascal_iter_20000.caffemodel',
        'symmax300': '/disk/liuzhe/PyCharmProjects/skeleton/snapshot/h2-symmax-iter20000.caffemodel',
        'wh-symmax': '/disk/liuzhe/PyCharmProjects/skeleton/snapshot/h2-wh-symmax_iter_40000.caffemodel',
    },
}


def deepflux(net):
    threshold = 0.4
    dks = 7
    eks = 9

    def forward(x, skeleton=True):
        x = x.astype(np.float32)
        x -= np.array((103.939, 116.779, 123.68))
        in_ = x.transpose((2, 0, 1))
        net.blobs['data'].reshape(1, *in_.shape)
        net.blobs['data'].data[...] = in_
        net.forward()

        flux = net.blobs['fcrop'].data
        skl = flux_to_skl(flux, threshold, dks, eks)

        return skl if skeleton else flux
    return forward


def fsds_hifi(net, output, ms):
    EPSILON = 1e-6

    def expand_channel(img):
        if img.ndim == 2:
            img = img[:, :, np.newaxis]
            img = np.repeat(img, 3, 2)
        return img


    def getsk(x):
        x = np.squeeze(x)
        if x.ndim == 2:
            return x
        elif x.ndim == 3:
            return 1 - np.squeeze(x[0, :, :])
        else:
            raise Exception("Invalid blob ndim: %d"%x.ndim)


    def forward(img):
        assert img.ndim == 3, "image shape: %s"%str(img.shape)
        img -= np.array((104.00698793,116.66876762,122.67891434))
        img = img.transpose((2, 0, 1))
        net.blobs['data'].reshape(1, *img.shape)
        net.blobs['data'].data[...] = img
        return net.forward()


    def main(x):
        x = x.astype(np.float32)
        x = expand_channel(x)
        h, w, _ = x.shape
        skeleton = np.zeros([h, w], np.float32)
        
        if ms:
            scales = np.array([0.5, 1, 1.5])
        else:
            scales = np.array([1])

        for s in scales:
            h1, w1 = int(s*h), int(s*w)
            img1 = cv2.resize(x, (w1, h1), interpolation=cv2.INTER_CUBIC).astype(np.float32)
            result = forward(img1)
            sk1 = getsk(result[output])
            sk1 = cv2.resize(sk1, (w, h), interpolation=cv2.INTER_CUBIC).astype(np.float32)
            skeleton += sk1
        skeleton /= len(scales)
        skeleton /= max(skeleton.max(), EPSILON)
        return skeleton
    return main


# net_pools = {}
def load_model(name, dataset, cuda):
    # model_key = name + dataset + str(cuda)
    # if model_key in net_pools:
    #     net = net_pools[model_key]
    # else:
    if name == 'hifi' and dataset == 'sklarge':
        proto = PROTOCOL['hifi-sklarge']
    else:
        proto = PROTOCOL[name]
    model_path = MODEL_PATH[name][dataset]

    if name == 'fsds' or name == 'hifi':
        caffe_fsds.set_device(int(cuda))
        caffe_fsds.set_mode_gpu()
        net = caffe_fsds.Net(proto, model_path, caffe_fsds.TEST)
    else:
        caffe.set_device(int(cuda))
        caffe.set_mode_gpu()
        net = caffe.Net(proto, model_path, caffe.TEST)
    # net_pools[model_key] = net

    if name == 'deepflux':
        return deepflux(net)
    elif name == 'fsds':
        return fsds_hifi(net, 'softmax_fuse', ms=False)
    elif name == 'hifi':
        return fsds_hifi(net, 'softmax_h1h2_fuse', ms=True)
    else:
        return False
